from flask import Flask, request, jsonify
from flask_restful import Resource, Api
from zhipuai import ZhipuAI

class ChatAPI(Resource):
    def post(self):
        data = request.json
        api_key = data.get('api_key')
        messages = data.get('messages')
        model_name = data.get('model_name')
        if not api_key or not messages:
            return {"message": "api_key and messages are required"}, 400

        # 创建 ZhipuAI 客户端
        client = ZhipuAI(api_key=api_key)

        try:
            # 调用 chat completions
            response = client.chat.completions.create(
                model=model_name,  # 模型名称
                messages=messages
            )

            # 格式化返回结果
            response_data = [{"role": choice.message.role, "content": choice.message.content} for choice in response.choices]

            return jsonify(response_data)

        except Exception as e:
            return {"message": str(e)}, 500